# -*- coding:utf-8 -*-
import requests
from pprint import pprint
from threading import Lock
import json
response = requests.get('http://10.31.3.154:1025/v1/models')
data = response.json()
model_name = data['data'][0]['id'] if data['data'] else None

class model_infer:
    def __init__(self):
        self.state = {"code":200, "isEnd":False, "message":""}
        self.init_client()
        self.lock = Lock()
        self.memory_list = [{'role': 'system', 'content': '你是一个优秀的生活助理。' }]

    def init_client(self):
        self.url = "http://10.31.3.154:1025/v1/chat/completions"
        self.headers = {"Content-Type": "application/json"}

    def getState(self,):
        with self.lock:
            return self.state.copy()

    def predict(self, text):
        '''infer '''
        with self.lock:
            self.state['isEnd'], self.state['message'] = False, ""
        if text == "":
            return
            

        self.memory_list.append(
            {'role': 'user', 'content': text }
        )

        data = {
            "model":model_name,
            "messages": self.memory_list,
            "stream": False,
            "presence_penalty": 1.03,
            "frequency_penalty": 1.0,
            "repetition_penalty": 1.0,
            "temperature": 0.5,
            "top_p": 0.95,
            "top_k": -1,
            "seed": 1,
            "stop": ["stop1", "stop2"],
            "stop_token_ids": [2],
            "ignore_eos": False,
            "max_tokens": 100,
            "tools": [{
                    "type": "function",
                    "function": {
                        "name": "get_delivery_date",
                        "strict": True,
                        "description": "每当您服务客户的时候，你擅长查询各类信息",
                        "parameters": {
                            "type": "object",
                            "properties": {
                                "order_id": {
                                "type": "string",
                                "description": "当然可以。"
                                }
                            },
                            "required": [
                                "order_id"
                            ],
                            "additionalProperties": False
                        }
                    }
                }
            ],
        "tool_choice": "auto",
        }
        response = requests.post(self.url, headers=self.headers, data=json.dumps(data))
        data = dict(response.json())

        try:
            origin_message = data['choices'][0]['message']
            new_message = {
                'role': origin_message['role'],
                'content': origin_message['content']
            }
            self.memory_list.append(new_message)
            with self.lock:
                self.state['message'] = new_message['content']
        except:        
            pprint(f'{data=}')
            with self.lock:
                self.state['message'] = str(data)
        
        self.state['isEnd'] = True
        return self.state['message']

    def predict_stream(self, text):
        '''TODO'''
        with self.lock:
            self.state['isEnd'], self.state['message'] = False, ""
        if text == "":
            return
            

        self.memory_list.append(
            {'role': 'user', 'content': text }
        )

        data = {
            "model": model_name,
            "messages": self.memory_list,
            "stream": True,
            "presence_penalty": 1.03,
            "frequency_penalty": 1.0,
            "repetition_penalty": 1.0,
            "temperature": 0.5,
            "top_p": 0.95,
            "top_k": -1,
            "seed": 1,
            "stop": ["stop1", "stop2"],
            "stop_token_ids": [2],
            "ignore_eos": False,
            "max_tokens": 1024,
            "tools": [{
                    "type": "function",
                    "function": {
                        "name": "get_delivery_date",
                        "strict": True,
                        "description": "每当您服务客户的时候，你擅长查询各类信息",
                        "parameters": {
                            "type": "object",
                            "properties": {
                                "order_id": {
                                "type": "string",
                                "description": "当然可以。"
                                }
                            },
                            "required": [
                                "order_id"
                            ],
                            "additionalProperties": False
                        }
                    }
                }
            ],
        "tool_choice": "auto",
        }
        response = requests.post(self.url, 
                                 headers={"Content-Type": "application/json",
                                                    "Cache-Control":"no-cache",
                                                    "Connection":"keep-alive"}, 
                                data=json.dumps(data),stream=True)
        
        full_s = ""
        print('Answer:',end="")
        for s in response.iter_lines():
            if s not in [b'',b'data: [DONE]']:
                token = dict(json.loads(s.decode().strip("data:").strip()))
                temp_s = token['choices'][0]['delta']['content']
                full_s += temp_s
                print(temp_s,end="")
                with self.lock:
                    self.state['message'] = full_s
        print("")
        self.state['isEnd'] = True
        
        new_message = {
            'role': 'assistant',
            'content': full_s
        }
        self.memory_list.append(new_message)


if __name__ == "__main__":
    infer_engine = model_infer()
    while True:
        line = input("Question:")
        infer_engine.predict_stream(line) #流式
        #print(f"Answer:{infer_engine.predict(line)}") #非流式
